'''
Created on Nov 23, 2009

@author: mkiyer
'''

from phdb import PileupDB, PileupQuery
from veggie.db.sample.samplegroup import parse_samplegroups_xml
import random
import logging

def get_bed_intervals(fhd, numlines):
    intervals = []
    for linenum, line in enumerate(fhd):
        if linenum >= numlines:
            break
        fields = line.strip().split('\t')
        chrom, start, end = fields[0], int(fields[1]), int(fields[2])
        intervals.append((chrom, start, end))
    return intervals

def get_rand_intervals(n, interval_length):
    rands = [random.randint(0, 240000000) for x in xrange(n)]
    intervals = [('chr1', x, x + interval_length) for x in rands]
    return intervals

def get_sample_library_map(sample_pool_file):
    pools = {}
    for line in open(sample_pool_file):
        fields = line.strip().split('\t')
        sample_name = fields[0]
        library_ids = fields[1].split(',')
        pools[sample_name] = library_ids
    return pools

if __name__ == '__main__':
    logging.basicConfig(level=logging.DEBUG)
    import sys
    import numpy as np
    from scipy.stats.stats import pearsonr
    bed_file = sys.argv[1]
    phdb_path = sys.argv[2]
    sample_pools_path = sys.argv[3]
    sample_library_map = get_sample_library_map(sample_pools_path)

    #sample_groups = list(parse_samplegroups_xml(sys.argv[1]))
    #samples = []
    #for sgroup in sample_groups:
    #    samples.extend(sgroup.samples)
    #samples = ['VCaP', 'LNCaP']
    samples = ['VCaP', 'aT8_2', 'aN13_2']
    print 'Samples:', samples

    # load query
    myquery2 = PileupQuery(samples, phdb_path, sample_library_map)
    print 'new samples:', myquery2.index_sample_names, myquery2.sample_indexes_map
    nlibraries = len(myquery2.h5groups)
    nsamples = len(myquery2.ordered_sample_names)
    
    # get some intervals
    num_intervals = 1000
    intervals = get_bed_intervals(open(bed_file), num_intervals)
    #intervals = get_rand_intervals(100, 10000)
    
    print 'querying'
    batch_sample2 = np.memmap('outs.ndarray', dtype=np.float, mode="w+", shape=(num_intervals, nsamples))
    myquery2.batch_sample_coverage2(intervals, batch_sample2, norm=True)

#    sample_result = np.memmap('out4.ndarray', dtype=np.float, mode="w+", shape=(num_intervals, nsamples))
#    for i, interval in enumerate(intervals):
#        chrom, start, end = interval
#        arr2 = myquery2.sample_coverage(chrom, start, end, norm=True)
#        sample_result[i,:] = myquery2.to_rpkm_2d(arr2)
#
#    print sample_result
    sys.exit(0)

    # create an array to store output

    batch_result = np.memmap('out1.ndarray', dtype=np.float, mode="w+", shape=(num_intervals, nlibraries))
    myquery2.batch_coverage(intervals, batch_result, norm=True)
    batch_sample = np.memmap('outs.ndarray', dtype=np.float, mode="w+", shape=(num_intervals, nsamples))    
    myquery2.batch_sample_coverage(intervals, batch_sample, norm=True)
    batch_sample2 = np.memmap('outs.ndarray', dtype=np.float, mode="w+", shape=(num_intervals, nsamples))    
    myquery2.batch_sample_coverage2(intervals, batch_sample2, norm=True)

    # do the same query in normal style
    result = np.memmap('out2.ndarray', dtype=np.float, mode="w+", shape=(num_intervals, nlibraries))
    sample_result = np.memmap('out4.ndarray', dtype=np.float, mode="w+", shape=(num_intervals, nsamples))    

    for i, interval in enumerate(intervals):
        chrom, start, end = interval
        arr = myquery2.coverage(chrom, start, end, norm=True)
        result[i,:] = myquery2.to_rpkm_2d(arr)
        arr2 = myquery2.sample_coverage(chrom, start, end, norm=True)
        sample_result[i,:] = myquery2.to_rpkm_2d(arr2)
        print result[i,:], batch_result[i,:], sample_result[i,:], batch_sample[i,:], batch_sample2[i,:]
#    print result
#    print sample_result

    sys.exit(0)


    print chrom, start, end
    arr1 = myquery2.coverage(chrom, start, end, norm=True)
    arr2 = myquery2.sample_coverage(chrom, start, end, True)
    rpkm1 = myquery2.to_rpkm_2d(arr1)
    rpkm2 = myquery2.to_rpkm_2d(arr2)
    print 'arr1', arr1
    print 'arr2', arr2
    print 'rpkm1', rpkm1
    print 'rpkm2', rpkm2
    sys.exit(0)

    ntests = 100
    corrs = []
    rands = [random.randint(0, 240000000) for x in xrange(ntests)]
    intervals = [('chr1', x, x+10) for x in rands]
    for interval in intervals:
        chrom, start, end = interval
#        rpkm1 = myquery.coverage(chrom, start, end, norm=True, rpkm=True)              
#        arr2 = myquery2.coverage(chrom, start, end, norm=True)
#        rpkm2 = myquery2.to_rpkm_2d(arr2)
        
        #rpkm1 = myquery.sample_coverage(chrom, start, end, True, True)
        #print 'old', rpkm1
        print chrom, start, end
        arr1 = myquery2.coverage(chrom, start, end, norm=True)
        arr2 = myquery2.sample_coverage(chrom, start, end, True)
        rpkm1 = myquery2.to_rpkm_2d(arr1)
        rpkm2 = myquery2.to_rpkm_2d(arr2)
        print 'arr1', arr1
        print 'arr2', arr2
        print 'rpkm1', rpkm1
        print 'rpkm2', rpkm2
        #r,p = pearsonr(arr1[0], arr2[0])
        #print r,p
        #corrs.append(r)
        
    #print np.min(corrs), np.max(corrs), np.median(corrs), np.average(corrs)
    sys.exit(0)
    intervals = []
    count = 0

    for line in open(bed_file):
        fields = line.strip().split('\t')
        chrom, start, end = fields[0], int(fields[1]), int(fields[2])
        arr1 = myquery.coverage(chrom, start, end, norm=True)                
        arr2 = myquery2.coverage(chrom, start, end, norm=True)
        print np.average(arr1), np.average(arr2)
        print arr1.shape, arr2.shape
        r,p = pearsonr(arr1[0], arr2[0])
        corrs.append(r)
        count += 1
        if count > 100:
            break
        #print arr1
        #print arr2
        intervals.append((chrom, start, end))
    
    import matplotlib.pyplot as plt

    sys.exit(0)
    
    rands = [random.randint(0, 240000000) for x in xrange(100)]
    intervals = [('chr1', x, x+10000) for x in rands]
    myquery2.batch_coverage(intervals, None, norm=True)

    sys.exit(0)

    print 'Starting random queries on chromosome 1...'   
    for i in xrange(100):
        start = random.randint(0, 240000000)
        end = start + 10000
        print i, start, end
        arr1 = myquery.coverage('chr1', start, end, norm=True)                
        arr2 = myquery2.coverage('chr1', start, end, norm=True)                
        print arr1.shape
        print arr2.shape        
        print np.array_equal(arr1, arr2)

    phdb.close()
    myquery2.close()
